
import io.cloudsoft.winrm4j.client.ShellCommand;
import io.cloudsoft.winrm4j.client.WinRm;
import io.cloudsoft.winrm4j.client.WinRmClient;
import io.cloudsoft.winrm4j.client.WinRmClientBuilder;
import io.cloudsoft.winrm4j.client.WinRmClientContext;
import io.cloudsoft.winrm4j.client.shell.CommandStateType;
import io.cloudsoft.winrm4j.client.shell.ReceiveResponse;
import io.cloudsoft.winrm4j.client.shell.StreamType;
import io.cloudsoft.winrm4j.client.wsman.OptionSetType;
import io.cloudsoft.winrm4j.client.wsman.OptionType;
import io.cloudsoft.winrm4j.winrm.WinRmTool;
import io.cloudsoft.winrm4j.winrm.WinRmToolResponse;
import org.apache.commons.collections4.CollectionUtils;
import org.apache.commons.lang3.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.xml.bind.DatatypeConverter;
import javax.xml.ws.BindingProvider;
import java.io.StringWriter;
import java.io.UnsupportedEncodingException;
import java.lang.reflect.Field;
import java.lang.reflect.InvocationHandler;
import java.lang.reflect.Method;
import java.lang.reflect.Proxy;
import java.math.BigInteger;
import java.nio.charset.StandardCharsets;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

/**
 * WinRm 协议命令执行器
 * <a href="https://blog.csdn.net/lnktoking/article/details/110286621">使用说明</a>
 * 依赖包版本: io.cloudsoft.windows:winrm4j:0.10.0; org.apache.commons:commons-lang3:3.10
 * 原winrm4j的github: https://github.com/cloudsoft/winrm4j
 *
 * @Author linnankun
 * @Date 2020/11/28
 */
public class WinRmExecutor implements AutoCloseable {
    private final static Logger logger = LoggerFactory.getLogger(WinRmExecutor.class);

    private final static String[] escapeFrom;
    private final static String[] escapeTo;
    private final static Pattern matchPort = Pattern.compile(".*:(\\d+)$");

    private final WinRmClientContext context;
    private WinRmClientBuilder winRmClientBuilder;

    static {
        Map<String, String> escapeMap = new HashMap<String, String>() {{
            put("^", "^^");
            put("<", "^<");
            put(">", "^>");
            put("|", "^|");
            put("&", "^&");
//            put("%", "%%");
//            put("\\", "\\\\");
        }};

        escapeFrom = escapeMap.keySet().toArray(new String[0]);
        escapeTo = escapeMap.values().toArray(new String[0]);
    }

    public static void main(String[] args) {
        final WinRmClientBuilder builder = WinRmClient
                .builder(WinRmExecutor.getEndpointUrl("127.0.0.1", null))
//                .builder("http://192.168.7.70:5985/wsman")
                .credentials("test", "123456");
        try (WinRmExecutor executor = new WinRmExecutor(builder)) {
//            executor.setUseHttps(true);
            WinRmToolResponse response = executor.executeCmd("ipconfig");
            System.out.println(response.getStatusCode());
            System.out.println(response.getStdOut());
            System.out.println(response.getStdErr());

            System.out.println("---------------------");
            response = executor.executePs("echo $env:temp");
            System.out.println(response.getStatusCode());
            System.out.println(response.getStdOut());
            System.out.println(response.getStdErr());

            System.out.println("---------------------");
            String script = "@echo off\n" +
                    "echo %1 %2\n" +
                    "echo %3";
            response = executor.executeBatScript(script, "aa bb 123", null);
            System.out.println(response.getStatusCode());
            System.out.println(response.getStdOut());
            System.out.println(response.getStdErr());
        }
    }

    /**
     * WinRm 协议执行器
     */
    public WinRmExecutor(WinRmClientBuilder builder) {
        context = WinRmClientContext.newInstance();
        this.winRmClientBuilder = builder;
        builder.locale(Locale.CHINESE);
    }

    /**
     * 执行命令行
     */
    public WinRmToolResponse executeCmd(String command) {

        StringWriter out = new StringWriter();
        StringWriter err = new StringWriter();
        try (WinRmClient winRmClient = winRmClientBuilder.build()) {
            proxyWinRm(winRmClient);
            try (ShellCommand shell = winRmClient.createShell()) {
                int code = shell.execute(command, out, err);
                return new WinRmToolResponse(out.toString(), err.toString(), code);
            }
        }
    }

    /**
     * 执行bat脚本
     *
     * @param scriptContent 脚本内容
     * @param executeParams 执行参数
     * @param fileName      生成临时脚本时的脚本名，可空
     */
    public WinRmToolResponse executeBatScript(String scriptContent, String executeParams, String fileName) {
        boolean psScript = false;
        if (StringUtils.isBlank(fileName)) {
            fileName = "temp.bat";
        } else if (StringUtils.endsWithIgnoreCase(fileName, ".ps1")) {
            psScript = true;
        } else if (!StringUtils.endsWithIgnoreCase(fileName, ".bat")) {
            fileName += ".bat";
        }

        String scriptPath = null;
        try {
            scriptPath = writeTempScriptFile(scriptContent, fileName);
            String cmd = scriptPath;
            if (StringUtils.isNotBlank(executeParams)) {
                cmd += " " + executeParams;
            }
            if (psScript) {
                return executePs(cmd);
            } else {
                return executeCmd(cmd);
            }
        } finally {
            if (scriptPath != null) {
                final WinRmToolResponse response = deleteFile(scriptPath);
                if (response.getStatusCode() != 0) {
                    logger.warn("通过winrm执行脚本结束，但删除临时脚本失败：{}",
                            StringUtils.defaultString(response.getStdErr(), response.getStdOut()));
                }
            }
        }
    }


    /**
     * 执行power shell命令
     */
    public WinRmToolResponse executePs(String command) {
        byte[] bytes = command.getBytes(StandardCharsets.UTF_16LE);
        final String cmd = DatatypeConverter.printBase64Binary(bytes);
        return executeCmd("powershell -encodedcommand " + cmd);
    }


    /**
     * 执行power shell脚本
     *
     * @param scriptContent 脚本内容
     * @param executeParams 执行参数
     * @param fileName      生成临时脚本时的脚本名，可空
     */
    public WinRmToolResponse executePsScript(String scriptContent, String executeParams, String fileName) {
        if (StringUtils.isBlank(fileName)) {
            fileName = "temp.ps1";
        } else if (!StringUtils.endsWithIgnoreCase(fileName, ".ps1")) {
            fileName += ".ps1";
        }
        return executeBatScript(scriptContent, executeParams, fileName);
    }

    /**
     * 生成临时文件
     */
    public String writeTempScriptFile(String scriptContent, String fileName) {
        if (StringUtils.isBlank(scriptContent)) {
            throw new IllegalArgumentException("脚本内容不能为空");
        }
        if (StringUtils.isBlank(fileName)) {
            throw new IllegalArgumentException("文件名不能为空");
        }

        // 获取用户临时目录
        WinRmToolResponse response = executeCmd("echo %temp%");
        if (response.getStatusCode() != 0) {
            throw new RuntimeException(String.format("生成临时执行脚本失败[exitCode:%s]：%s", response.getStatusCode(),
                    StringUtils.defaultString(response.getStdErr(), response.getStdOut())));
        }
        // 将脚本写入临时目录
        String tempFilePath = response.getStdOut().trim() + "\\" + System.currentTimeMillis() + "_" + fileName;
        String[] lines = scriptContent.split("\n");
        StringBuilder writeScriptCmd = new StringBuilder();
        for (int i = 0; i < lines.length; i++) {
            String outChar = " >> ";
            if (i == 0) {
                outChar = " > ";
            } else {
                writeScriptCmd.append(" & ");
            }

            if (StringUtils.isBlank(lines[i])) {
                writeScriptCmd.append("echo=").append(outChar).append(tempFilePath);
            } else {
                writeScriptCmd.append("echo ").append(escapeSpecialChar(lines[i])).append(outChar).append(tempFilePath);
            }
        }

        response = executeCmd(writeScriptCmd.toString());
        if (response.getStatusCode() != 0) {
            throw new RuntimeException(String.format("写入临时执行脚本失败[exitCode:%s]：%s", response.getStatusCode(),
                    StringUtils.defaultString(response.getStdErr(), response.getStdOut())));
        }
        return tempFilePath;
    }

    /**
     * 删除文件
     */
    private WinRmToolResponse deleteFile(String filePath) {
        String deleteCmd = "del /Q \"" + filePath + "\"";
        return executeCmd(deleteCmd);
    }

    private String escapeSpecialChar(String text) {
        return StringUtils.replaceEach(text, escapeFrom, escapeTo);
    }

    public static String getEndpointUrl(String address, Boolean useHttps) {
        if (address.startsWith("http:") || address.startsWith("https:")) {
            if (useHttps != null) {
                if (useHttps && address.startsWith("http:"))
                    throw new IllegalArgumentException("Invalid setting useHttps and address starting http://");
                if (!useHttps && address.startsWith("https:"))
                    throw new IllegalArgumentException("Invalid setting useHttp and address starting https://");
            }
            return address;
        } else {
            Matcher matcher = matchPort.matcher(address);

            if (matcher.matches()) {
                if (useHttps == null) {
                    useHttps = matcher.group(1).equals("5986");
                }
                return (useHttps ? "https" : "http") + "://" + address + "/wsman";
            } else {
                if (useHttps != null && useHttps) {
                    return "https://" + address + ":" + WinRmTool.DEFAULT_WINRM_HTTPS_PORT + "/wsman";
                } else {
                    return "http://" + address + ":" + WinRmTool.DEFAULT_WINRM_PORT + "/wsman";
                }
            }
        }
    }

    /**
     * 代理 WinRmClient 的 winrm 字段，修改解析编码
     */
    private static void proxyWinRm(WinRmClient winRmClient) {
        try {
            final Field winrmField = WinRmClient.class.getDeclaredField("winrm");
            winrmField.setAccessible(true);
            final Object winrm = winrmField.get(winRmClient);
            Object proxyWinrm = Proxy.newProxyInstance(WinRm.class.getClassLoader(),
                    new Class[]{WinRm.class, BindingProvider.class},
                    new InvocationHandler() {
                        @Override
                        public Object invoke(Object proxy, Method method, Object[] args) throws Throwable {
                            if ("create".equals(method.getName())) {
                                OptionSetType optionSetType = (OptionSetType) args[5];
                                for (OptionType optionType : optionSetType.getOption()) {
                                    // 设置执行命令使用GBK字符编码
                                    if ("WINRS_CODEPAGE".equals(optionType.getName())) {
                                        optionType.setValue("936");
                                    }
                                }
                            }
                            Object value = method.invoke(winrm, args);
                            if ("receive".equals(method.getName())) {
                                value = new ReceiveResponseProxy((ReceiveResponse) value);
                            }
                            return value;
                        }
                    });
            winrmField.set(winRmClient, proxyWinrm);
        } catch (Exception e) {
            logger.error("设置编码失败", e);
        }
    }

    @Override
    public void close() {
        context.shutdown();
    }

    private static class ReceiveResponseProxy extends ReceiveResponse {
        private ReceiveResponse target;

        private ReceiveResponseProxy(ReceiveResponse response) {
            this.target = response;
        }

        @Override
        public List<StreamType> getStream() {
            List<StreamType> list = target.getStream();
            if (CollectionUtils.isNotEmpty(list)) {
                for (StreamType streamType : list) {
                    byte[] bytes = streamType.getValue();
                    if (bytes == null || bytes.length == 0) {
                        continue;
                    }

                    try {
                        // 接收的字符是GBK编码，转为系统默认编码
                        streamType.setValue(new String(bytes, "GBK").getBytes());
                    } catch (UnsupportedEncodingException e) {
                        logger.error("使用GBK解析结果失败", e);
                    }
                }
            }
            return list;
        }

        @Override
        public CommandStateType getCommandState() {
            return target.getCommandState();
        }

        @Override
        public void setCommandState(CommandStateType value) {
            target.setCommandState(value);
        }

        @Override
        public BigInteger getSequenceID() {
            return target.getSequenceID();
        }

        @Override
        public void setSequenceID(BigInteger value) {
            target.setSequenceID(value);
        }
    }

}
